massive changes / restucturing

app.py

@@ -1,33 +1,23 @@
-import spaces
 import gradio as gr
-import torch
-from diffusers import DiffusionPipeline
-import rembg
-
-pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
-pipe.to("cuda")
-
-# Function to generate an image from text using diffusion
-@spaces.GPU
-def generate_image(prompt):
-    prompt += "no background, side view, minimalist shot"
-    
-    image = pipe(prompt).images[0]
-    image2 = rembg.remove(image)
-
-    return image, image2
+from util.text_img import generate_image
 
 _TITLE = "Shoe Generator"
 with gr.Blocks(_TITLE) as ShoeGen:
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Enter a discription of a shoe")
-            # neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
-            button_gen = gr.Button("Generate Image")
-        with gr.Column():
-            image = gr.Image(label="Generated Image", show_download_button=True) 
-            image_nobg = gr.Image(label="Generated Image (No Background)", show_download_button=True) 
+    with gr.Tab("Text to Image Generator"):
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(label="Enter a discription of a shoe")
+                # neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
+                button_gen = gr.Button("Generate Image")
+            with gr.Column():
+                with gr.Tab("With Background"):
+                    image = gr.Image(label="Generated Image", show_download_button=True, show_label=False)
+                with gr.Tab("Without Background"): 
+                    image_nobg = gr.Image(label="Generated Image", show_download_button=True, show_label=False) 
     
-    button_gen.click(generate_image, inputs=[prompt], outputs=[image, image_nobg])
+        button_gen.click(generate_image, inputs=[prompt], outputs=[image, image_nobg])
 
+    with gr.Tab("Image to Video Generator (V3D)"):
+        pass
+    
 ShoeGen.launch()

util/text_img.py

@@ -0,0 +1,17 @@
+import spaces
+import rembg
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+pipe.to("cuda")
+
+# Function to generate an image from text using diffusion
+@spaces.GPU
+def generate_image(prompt):
+    prompt += "no background, side view, minimalist shot"
+    
+    image = pipe(prompt).images[0]
+    image2 = rembg.remove(image)
+
+    return image, image2