Update app.py

app.py

@@ -50,67 +50,57 @@ pipeline.set_ip_adapter_scale([0.7, 0.3])
 pipeline.enable_model_cpu_offload()
 pipeline.enable_vae_tiling()
 
-def generate_storybook(data):
-    print("📥 Input JSON received:")
+def generate_single_scene(data):
+    print("📥 Single scene input received:")
     print(json.dumps(data, indent=2))
 
-    character_image_url = data["character_image_url"]
-    style = data["style"]
-    scenes = data["scenes"]
-
-    face_image = load_image(character_image_url)
-    style_images = [load_image(url) for url in STYLE_MAP.get(style, [])]
-
-    images = []
-
-    for i, prompt in enumerate(scenes):
-        print(f"🎬 Generating scene {i+1}: {prompt}")
-        try:
-            torch.cuda.empty_cache()
-            gc.collect()
-
-            result = pipeline(
-                prompt=prompt,
-                ip_adapter_image=[style_images, face_image],
-                negative_prompt="blurry, bad anatomy, low quality",
-                width=448,
-                height=672,
-                guidance_scale=5.0,
-                num_inference_steps=15,
-                generator=torch.Generator(device).manual_seed(i + 42)
-            )
-
-            image = result.images[0] if hasattr(result, "images") else result
-            print(f"🖼️ Image type: {type(image)}")
-
-            if isinstance(image, Image.Image):
-                images.append(image)
-                print(f"✅ Scene {i+1} added to image list.")
-            else:
-                print(f"⚠️ Scene {i+1} is not a valid image object.")
-
-        except Exception as e:
-            print(f"❌ Exception during scene {i+1}: {e}")
-            traceback.print_exc()
-
-    print(f"📦 Returning {len(images)} image(s)")
-    return images
-
-def generate_storybook_from_textbox(json_input_text):
+    try:
+        character_image_url = data["character_image_url"]
+        style = data["style"]
+        scene_prompt = data["scene"]
+
+        face_image = load_image(character_image_url)
+        style_images = [load_image(url) for url in STYLE_MAP.get(style, [])]
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        result = pipeline(
+            prompt=scene_prompt,
+            ip_adapter_image=[style_images, face_image],
+            negative_prompt="blurry, bad anatomy, low quality",
+            width=512,
+            height=768,
+            guidance_scale=5.0,
+            num_inference_steps=15,
+            generator=torch.Generator(device).manual_seed(42)
+        )
+
+        image = result.images[0] if hasattr(result, "images") else result
+        print(f"✅ Scene generated. Type: {type(image)}")
+
+        return image
+
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        traceback.print_exc()
+        return None
+
+def generate_from_json(json_input_text):
     try:
         data = json.loads(json_input_text)
-        return generate_storybook(data)
+        return generate_single_scene(data)
     except Exception as e:
         print(f"❌ JSON parse or generation error: {e}")
         traceback.print_exc()
-        return [f"Error: {str(e)}"]
+        return None
 
 iface = gr.Interface(
-    fn=generate_storybook_from_textbox,
-    inputs=gr.Textbox(label="Input JSON", lines=20, placeholder="{...}"),
-    outputs=gr.Gallery(label="Generated Story Scenes", show_label=True, columns=1),
-    title="AI Storybook Generator (Low VRAM Mode)",
-    description="Optimized for lower VRAM GPUs. Paste JSON to generate consistent scenes."
+    fn=generate_from_json,
+    inputs=gr.Textbox(label="Input JSON", lines=10, placeholder='{"character_image_url": "...", "style": "pixar", "scene": "..."}'),
+    outputs=gr.Image(label="Generated Scene"),
+    title="Single-Scene Storybook Generator",
+    description="Send one scene at a time to generate consistent character-based images."
 )
 
 iface.launch()