test gradio

app.py

@@ -1,20 +1,19 @@
 import gradio as gr
 import torch
-from PIL import Image
 from models.transformer_sd3 import SD3Transformer2DModel
 from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
 import os
+import spaces
 from huggingface_hub import login
-
 token = os.getenv("HF_TOKEN")
 login(token=token)
 
-# Model and paths
+# Model and Pipeline Setup
 model_path = 'stabilityai/stable-diffusion-3.5-large'
 ip_adapter_path = './ip-adapter.bin'
 image_encoder_path = "google/siglip-so400m-patch14-384"
 
-# Load SD3.5 pipeline and components
+# Load transformer and pipeline
 transformer = SD3Transformer2DModel.from_pretrained(
     model_path, subfolder="transformer", torch_dtype=torch.bfloat16
 )
@@ -22,6 +21,7 @@ pipe = StableDiffusion3Pipeline.from_pretrained(
     model_path, transformer=transformer, torch_dtype=torch.bfloat16
 ).to("cuda")
 
+# Initialize IP Adapter
 pipe.init_ipadapter(
     ip_adapter_path=ip_adapter_path,
     image_encoder_path=image_encoder_path,
@@ -29,38 +29,44 @@ pipe.init_ipadapter(
 )
 
 
-@gr.Interface()
-def gui_generation(image: Image, style_image: Image):
+@spaces.GPU
+def gui_generation(text, num_imgs, width, height):
     """
-    Generate an image based on input and style images.
+    Generate images using Stable Diffusion 3.5
     """
-    generator = torch.Generator("cuda").manual_seed(42)  # Reproducibility
-
-    output = pipe(
-        width=1024,
-        height=1024,
-        prompt="",
-        negative_prompt="",
+    images = pipe(
+        prompt=text,
+        width=width,
+        height=height,
+        num_images_per_prompt=num_imgs,
+        negative_prompt="lowres, low quality, worst quality",
         num_inference_steps=24,
         guidance_scale=5.0,
-        generator=generator,
-        clip_image=style_image,
-        ipadapter_scale=0.5,
-    ).images[0]
-    return output
+        generator=torch.Generator("cuda").manual_seed(42),
+    ).images
 
+    return images
 
-# Gradio UI elements
-image_input = gr.Image(type="pil", label="Input Image")
-style_image_input = gr.Image(type="pil", label="Style Image")
-output_image = gr.Image(label="Generated Image")
 
-interface = gr.Interface(
-    gui_generation,
-    inputs=[image_input, style_image_input],
-    outputs=output_image,
-    title="Image Generation with Style Image",
-    description="Upload an input image and a style image to generate a new image based on the style."
-)
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Stable Diffusion 3.5 Image Generation")
+
+    with gr.Row():
+        prompt_box = gr.Textbox(label="Prompt", placeholder="Enter your image generation prompt")
+        number_slider = gr.Slider(1, 30, value=2, step=1, label="Batch size")
+
+    with gr.Row():
+        width_slider = gr.Slider(256, 1536, value=1024, step=64, label="Width")
+        height_slider = gr.Slider(256, 1536, value=1024, step=64, label="Height")
+
+    gallery = gr.Gallery(columns=[3], rows=[1], object_fit="contain", height="auto")
+
+    generate_btn = gr.Button("Generate")
 
-interface.launch()
+    generate_btn.click(
+        fn=gui_generation,
+        inputs=[prompt_box, number_slider, width_slider, height_slider],
+        outputs=gallery
+    )
+demo.launch()