Update app.py

app.py

@@ -1,138 +1,88 @@
-import torch
 import gradio as gr
 import numpy as np
 import random
+import torch
+from diffusers import DiffusionPipeline
 
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionInstructPix2PixPipeline,
-    StableVideoDiffusionPipeline,
-    WanPipeline,
-)
-from diffusers.utils import export_to_video, load_image
-
-# Force CPU mode
-device = "cpu"
-dtype = torch.float32
-MAX_SEED = np.iinfo(np.int32).max
-
-# Global pipeline holders
-TXT2IMG_PIPE = None
-IMG2IMG_PIPE = None
-TXT2VID_PIPE = None
-IMG2VID_PIPE = None
+# Define available models and their corresponding Hugging Face repositories
+MODEL_REPOS = {
+    "Stable Diffusion XL Base 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
+    "SDXL-Turbo": "stabilityai/sdxl-turbo",
+    "Playground v2 1024px Aesthetic": "playgroundai/playground-v2-1024px-aesthetic",
+    "Segmind Vega": "segmind/Segmind-Vega",
+    "SSD-1B": "segmind/SSD-1B",
+    "Kandinsky 3": "kandinsky-community/kandinsky-3",
+    "PixArt-LCM-XL-2-1024-MS": "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
+    "BLIP Diffusion": "salesforce/blipdiffusion",
+    "Muse-512-Finetuned": "amused/muse-512-finetuned",
+    "Flux 1 Dev": "black-forest-labs/FLUX.1-dev"
+}
 
-# Helper to load models
-def make_pipe(cls, model_id, **kwargs):
-    pipe = cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
-    pipe.to(device)
-    return pipe
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Text → Image
-def generate_image_from_text(prompt, seed, randomize_seed):
-    global TXT2IMG_PIPE
-    if TXT2IMG_PIPE is None:
-        TXT2IMG_PIPE = make_pipe(StableDiffusionPipeline, "stabilityai/stable-diffusion-2-1-base")
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.manual_seed(seed)
-    image = TXT2IMG_PIPE(prompt=prompt, num_inference_steps=20, generator=generator).images[0]
-    return image, seed
+# Cache for loaded pipelines
+loaded_pipelines = {}
 
-# Image → Image
-def generate_image_from_image_and_prompt(image, prompt, seed, randomize_seed):
-    global IMG2IMG_PIPE
-    if IMG2IMG_PIPE is None:
-        IMG2IMG_PIPE = make_pipe(StableDiffusionInstructPix2PixPipeline, "timbrooks/instruct-pix2pix")
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.manual_seed(seed)
-    out = IMG2IMG_PIPE(prompt=prompt, image=image, num_inference_steps=8, generator=generator)
-    return out.images[0], seed
+# Maximum seed value
+MAX_SEED = np.iinfo(np.int32).max
 
-# Text → Video
-def generate_video_from_text(prompt, seed, randomize_seed):
-    global TXT2VID_PIPE
-    if TXT2VID_PIPE is None:
-        TXT2VID_PIPE = make_pipe(WanPipeline, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.manual_seed(seed)
-    frames = TXT2VID_PIPE(prompt=prompt, num_frames=12, generator=generator).frames[0]
-    return export_to_video(frames, "/tmp/wan_video.mp4", fps=8), seed
+def load_pipeline(model_name):
+    """Load and cache the pipeline for the selected model."""
+    if model_name in loaded_pipelines:
+        return loaded_pipelines[model_name]
+    repo_id = MODEL_REPOS[model_name]
+    try:
+        pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype)
+        pipeline.to(device)
+        loaded_pipelines[model_name] = pipeline
+        return pipeline
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model '{model_name}': {e}")
 
-# Image → Video
-def generate_video_from_image(image, seed, randomize_seed):
-    global IMG2VID_PIPE
-    if IMG2VID_PIPE is None:
-        IMG2VID_PIPE = make_pipe(
-            StableVideoDiffusionPipeline,
-            "stabilityai/stable-video-diffusion-img2vid-xt"
-        )
+def generate_image(prompt, model_name, width, height, guidance_scale, num_inference_steps, seed, randomize_seed):
+    """Generate an image using the selected model and parameters."""
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.manual_seed(seed)
-    image = load_image(image).resize((512, 288))
-    frames = IMG2VID_PIPE(image=image, num_inference_steps=16, generator=generator).frames[0]
-    return export_to_video(frames, "/tmp/svd_video.mp4", fps=8), seed
+    generator = torch.Generator(device=device).manual_seed(seed)
+    pipeline = load_pipeline(model_name)
+    try:
+        image = pipeline(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator
+        ).images[0]
+        return image, seed
+    except Exception as e:
+        raise RuntimeError(f"Image generation failed: {e}")
 
-# Gradio Interface
+# Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 AI Playground – Text & Image → Image & Video")
-
-    with gr.Tabs():
-        # Text → Image
-        with gr.Tab("Text → Image"):
-            prompt_txt = gr.Textbox(label="Prompt")
-            btn_txt2img = gr.Button("Generate")
-            result_img = gr.Image()
-            seed_txt = gr.Slider(0, MAX_SEED, value=42, label="Seed")
-            rand_txt = gr.Checkbox(label="Randomize seed", value=True)
-            btn_txt2img.click(
-                generate_image_from_text,
-                inputs=[prompt_txt, seed_txt, rand_txt],
-                outputs=[result_img, seed_txt]
-            )
-
-        # Image → Image
-        with gr.Tab("Image → Image"):
-            image_in = gr.Image(label="Input Image")
-            prompt_img = gr.Textbox(label="Edit Prompt")
-            btn_img2img = gr.Button("Generate")
-            result_img2 = gr.Image()
-            seed_img = gr.Slider(0, MAX_SEED, value=123, label="Seed")
-            rand_img = gr.Checkbox(label="Randomize seed", value=True)
-            btn_img2img.click(
-                generate_image_from_image_and_prompt,
-                inputs=[image_in, prompt_img, seed_img, rand_img],
-                outputs=[result_img2, seed_img]
-            )
-
-        # Text → Video
-        with gr.Tab("Text → Video"):
-            prompt_vid = gr.Textbox(label="Prompt")
-            btn_txt2vid = gr.Button("Generate")
-            result_vid = gr.Video()
-            seed_vid = gr.Slider(0, MAX_SEED, value=555, label="Seed")
-            rand_vid = gr.Checkbox(label="Randomize seed", value=True)
-            btn_txt2vid.click(
-                generate_video_from_text,
-                inputs=[prompt_vid, seed_vid, rand_vid],
-                outputs=[result_vid, seed_vid]
-            )
+    gr.Markdown("# 🖼️ Text-to-Image Generator with Multiple Models")
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
+            model_name = gr.Dropdown(label="Select Model", choices=list(MODEL_REPOS.keys()), value="Stable Diffusion XL Base 1.0")
+            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=512)
+            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=512)
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=7.5)
+            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=50)
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+            generate_button = gr.Button("Generate Image")
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+            output_seed = gr.Textbox(label="Used Seed", interactive=False)
 
-        # Image → Video
-        with gr.Tab("Image → Video"):
-            image_vid = gr.Image(label="Input Image")
-            btn_img2vid = gr.Button("Animate")
-            result_vid2 = gr.Video()
-            seed_vid2 = gr.Slider(0, MAX_SEED, value=999, label="Seed")
-            rand_vid2 = gr.Checkbox(label="Randomize seed", value=True)
-            btn_img2vid.click(
-                generate_video_from_image,
-                inputs=[image_vid, seed_vid2, rand_vid2],
-                outputs=[result_vid2, seed_vid2]
-            )
+    generate_button.click(
+        fn=generate_image,
+        inputs=[prompt, model_name, width, height, guidance_scale, num_inference_steps, seed, randomize_seed],
+        outputs=[output_image, output_seed]
+    )
 
-demo.queue()
-demo.launch(show_error=True)
+if __name__ == "__main__":
+    demo.launch()