Update app.py

app.py

@@ -8,106 +8,96 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 MAX_SEED = 2**32 - 1
 
-# --- Model lists ---
+# --- Model lists ordered by size (light to heavy) ---
 image_models = {
-    "Stable Diffusion 1.5": "runwayml/stable-diffusion-v1-5",
+    "Stable Diffusion 1.5 (light)": "runwayml/stable-diffusion-v1-5",
     "Stable Diffusion 2.1": "stabilityai/stable-diffusion-2-1",
-    "SDXL Base 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
+    "Dreamlike 2.0": "dreamlike-art/dreamlike-photoreal-2.0",
     "Playground v2": "playgroundai/playground-v2-1024px-aesthetic",
-    "Kandinsky 3": "kandinsky-community/kandinsky-3",
+    "Muse 512": "amused/muse-512-finetuned",
     "PixArt": "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
+    "Kandinsky 3": "kandinsky-community/kandinsky-3",
     "BLIP Diffusion": "Salesforce/blipdiffusion",
-    "Muse 512": "amused/muse-512-finetuned",
-    "Dreamlike 2.0": "dreamlike-art/dreamlike-photoreal-2.0",
-    "OpenJourney": "prompthero/openjourney"
-}
-
-video_models = {
-    "AnimateDiff": "animate-diff/animate-diff",
-    "CogVideoX-5b": "THUDM/CogVideoX-5b",
-    "HunyuanVideo": "tencent/HunyuanVideo",
-    "LTX-Video": "Lightricks/LTX-Video",
-    "ModelScope T2V": "damo-vilab/modelscope-text-to-video-synthesis",
-    "VideoCrafter": "videocrafter/videocrafter",
-    "Mochi-1": "mochi/mochi-1",
-    "Allegro": "allegro/allegro",
-    "OpenSora": "LanguageBind/Open-Sora-Plan-v1.2.0",
-    "Zer0Scope": "zero-scope/zero-scope"
+    "SDXL Base 1.0 (heavy)": "stabilityai/stable-diffusion-xl-base-1.0",
+    "OpenJourney (heavy)": "prompthero/openjourney"
 }
 
 text_models = {
-    "GPT-2": "gpt2",
+    "GPT-2 (light)": "gpt2",
     "GPT-Neo 1.3B": "EleutherAI/gpt-neo-1.3B",
-    "GPT-J 6B": "EleutherAI/gpt-j-6B",
     "BLOOM 1.1B": "bigscience/bloom-1b1",
+    "GPT-J 6B": "EleutherAI/gpt-j-6B",
     "Falcon 7B": "tiiuae/falcon-7b",
-    "MPT 7B": "mosaicml/mpt-7b",
-    "LLaMA 2 7B": "meta-llama/Llama-2-7b-hf",
-    "BTLM 3B": "cerebras/btlm-3b-8k-base",
     "XGen 7B": "Salesforce/xgen-7b-8k-base",
-    "StableLM 2": "stabilityai/stablelm-2-1_6b"
+    "BTLM 3B": "cerebras/btlm-3b-8k-base",
+    "MPT 7B": "mosaicml/mpt-7b",
+    "StableLM 2": "stabilityai/stablelm-2-1_6b",
+    "LLaMA 2 7B (heavy)": "meta-llama/Llama-2-7b-hf"
 }
 
-# --- Caching loaded pipelines ---
+# Cache
 image_pipes = {}
 text_pipes = {}
 
-# --- Functional logic ---
-def generate_image(prompt, model_name, seed, randomize_seed):
+def generate_image(prompt, model_name, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.manual_seed(seed)
 
+    progress(0, desc="Loading model...")
     if model_name not in image_pipes:
         image_pipes[model_name] = DiffusionPipeline.from_pretrained(
             image_models[model_name],
             torch_dtype=torch_dtype
         ).to(device)
-
     pipe = image_pipes[model_name]
-    image = pipe(prompt=prompt, generator=generator, num_inference_steps=25, width=512, height=512).images[0]
-    return image, seed
 
-def generate_text(prompt, model_name):
+    progress(25, desc="Running inference (step 1/3)...")
+    result = pipe(prompt=prompt, generator=generator, num_inference_steps=30, width=512, height=512)
+
+    progress(100, desc="Done.")
+    return result.images[0], seed
+
+def generate_text(prompt, model_name, progress=gr.Progress(track_tqdm=True)):
+    progress(0, desc="Loading model...")
     if model_name not in text_pipes:
         text_pipes[model_name] = pipeline("text-generation", model=text_models[model_name], device=0 if device == "cuda" else -1)
     pipe = text_pipes[model_name]
-    output = pipe(prompt, max_length=100, do_sample=True)[0]['generated_text']
-    return output
 
-def generate_video(prompt, model_name):
-    # Placeholder: real video models would return video frames
-    return f"[Video placeholder] Model: {model_name}\nPrompt: {prompt}"
+    progress(50, desc="Generating text...")
+    result = pipe(prompt, max_length=100, do_sample=True)[0]['generated_text']
+    progress(100, desc="Done.")
+    return result
 
-# --- Interface ---
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔄 Multi-Task AI Generator")
+    gr.Markdown("# 🧠 Multi-Model AI Playground with Progress")
 
     with gr.Tabs():
-        # Tab 1: Image Generation
-        with gr.Tab("🖼️ Image"):
+        # 🖼️ Image Gen Tab
+        with gr.Tab("🖼️ Image Generation"):
             img_prompt = gr.Textbox(label="Prompt")
-            img_model = gr.Dropdown(choices=list(image_models.keys()), value="Stable Diffusion 1.5", label="Select Image Model")
+            img_model = gr.Dropdown(choices=list(image_models.keys()), value="Stable Diffusion 1.5 (light)", label="Image Model")
             img_seed = gr.Slider(0, MAX_SEED, value=42, label="Seed")
             img_rand = gr.Checkbox(label="Randomize seed", value=True)
             img_btn = gr.Button("Generate Image")
             img_out = gr.Image()
             img_btn.click(fn=generate_image, inputs=[img_prompt, img_model, img_seed, img_rand], outputs=[img_out, img_seed])
 
-        # Tab 2: Video Generation
-        with gr.Tab("🎥 Video"):
-            vid_prompt = gr.Textbox(label="Prompt")
-            vid_model = gr.Dropdown(choices=list(video_models.keys()), value="AnimateDiff", label="Select Video Model")
-            vid_btn = gr.Button("Generate Video")
-            vid_out = gr.Textbox(label="Result (Placeholder)")
-            vid_btn.click(fn=generate_video, inputs=[vid_prompt, vid_model], outputs=vid_out)
-
-        # Tab 3: Text Generation
-        with gr.Tab("📝 Text"):
+        # 📝 Text Gen Tab
+        with gr.Tab("📝 Text Generation"):
             txt_prompt = gr.Textbox(label="Prompt")
-            txt_model = gr.Dropdown(choices=list(text_models.keys()), value="GPT-2", label="Select Text Model")
+            txt_model = gr.Dropdown(choices=list(text_models.keys()), value="GPT-2 (light)", label="Text Model")
             txt_btn = gr.Button("Generate Text")
-            txt_out = gr.Textbox(label="Generated Text")
+            txt_out = gr.Textbox(label="Output Text")
             txt_btn.click(fn=generate_text, inputs=[txt_prompt, txt_model], outputs=txt_out)
 
+        # 🎥 Video Gen Tab (placeholder)
+        with gr.Tab("🎥 Video Generation (Placeholder)"):
+            gr.Markdown("⚠️ Video generation is placeholder only. Models require special setup.")
+            vid_prompt = gr.Textbox(label="Prompt")
+            vid_btn = gr.Button("Pretend to Generate")
+            vid_out = gr.Textbox(label="Result")
+            vid_btn.click(lambda x: f"Fake video output for: {x}", inputs=[vid_prompt], outputs=[vid_out])
+
 demo.launch(show_error=True)