Update image_generator.py

image_generator.py

@@ -20,6 +20,13 @@ def flush_memory():
         torch.cuda.empty_cache()
 
 
+def get_device():
+    """Get the best available device."""
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
 def get_pipeline(model_id: str = "stabilityai/sdxl-turbo", pipeline_type: str = "t2i"):
     """
     Load and cache the diffusion pipeline.
@@ -31,25 +38,56 @@ def get_pipeline(model_id: str = "stabilityai/sdxl-turbo", pipeline_type: str =
     cache_key = f"{model_id}_{pipeline_type}"
 
     if cache_key not in _pipelines:
+        device = get_device()
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        
+        # Determine variant
+        variant = None
+        if device == "cuda" and ("turbo" in model_id or "sdxl" in model_id.lower()):
+            variant = "fp16"
+
         if pipeline_type == "t2i":
-            pipe = AutoPipelineForText2Image.from_pretrained(
-                model_id,
-                torch_dtype=torch.float16,
-                variant="fp16" if "turbo" in model_id or "sdxl" in model_id.lower() else None,
-            )
+            try:
+                pipe = AutoPipelineForText2Image.from_pretrained(
+                    model_id,
+                    torch_dtype=dtype,
+                    variant=variant,
+                )
+            except Exception:
+                # Fallback without variant
+                pipe = AutoPipelineForText2Image.from_pretrained(
+                    model_id,
+                    torch_dtype=dtype,
+                )
+            pipe = pipe.to(device)
             _pipelines[cache_key] = pipe
 
         elif pipeline_type == "i2i":
             # Try to reuse t2i components
             t2i_key = f"{model_id}_t2i"
             if t2i_key in _pipelines:
-                pipe = AutoPipelineForImage2Image.from_pipe(_pipelines[t2i_key])
+                try:
+                    pipe = AutoPipelineForImage2Image.from_pipe(_pipelines[t2i_key])
+                except Exception:
+                    try:
+                        pipe = AutoPipelineForImage2Image.from_pretrained(
+                            model_id, torch_dtype=dtype, variant=variant,
+                        )
+                    except Exception:
+                        pipe = AutoPipelineForImage2Image.from_pretrained(
+                            model_id, torch_dtype=dtype,
+                        )
+                    pipe = pipe.to(device)
             else:
-                pipe = AutoPipelineForImage2Image.from_pretrained(
-                    model_id,
-                    torch_dtype=torch.float16,
-                    variant="fp16" if "turbo" in model_id or "sdxl" in model_id.lower() else None,
-                )
+                try:
+                    pipe = AutoPipelineForImage2Image.from_pretrained(
+                        model_id, torch_dtype=dtype, variant=variant,
+                    )
+                except Exception:
+                    pipe = AutoPipelineForImage2Image.from_pretrained(
+                        model_id, torch_dtype=dtype,
+                    )
+                pipe = pipe.to(device)
             _pipelines[cache_key] = pipe
 
     return _pipelines[cache_key]
@@ -77,7 +115,6 @@ def get_model_config(model_id: str) -> Dict:
             "supports_guidance": True,
         },
     }
-    # Default config for unknown models
     return configs.get(model_id, {
         "num_inference_steps": 20,
         "guidance_scale": 7.0,
@@ -95,13 +132,7 @@ class ImageGenerator:
     def __init__(self, model_id: str = "stabilityai/sdxl-turbo"):
         self.model_id = model_id
         self.config = get_model_config(model_id)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    def _move_to_device(self, pipe):
-        """Move pipeline to the appropriate device."""
-        if self.device == "cuda":
-            pipe.to("cuda")
-        return pipe
+        self.device = get_device()
 
     def generate_first_panel(
         self,
@@ -115,22 +146,25 @@ class ImageGenerator:
     ) -> Image.Image:
         """Generate the first panel using text-to-image."""
         pipe = get_pipeline(self.model_id, "t2i")
-        pipe = self._move_to_device(pipe)
 
-        w, h = width or self.config["default_resolution"][0], height or self.config["default_resolution"][1]
+        w = width or self.config["default_resolution"][0]
+        h = height or self.config["default_resolution"][1]
         n_steps = steps or self.config["num_inference_steps"]
         cfg = guidance if guidance is not None else self.config["guidance_scale"]
         generator = torch.Generator(device=self.device).manual_seed(seed)
 
-        result = pipe(
+        kwargs = dict(
             prompt=prompt,
-            negative_prompt=negative_prompt if negative_prompt else None,
             num_inference_steps=n_steps,
             guidance_scale=cfg,
             width=w,
             height=h,
             generator=generator,
         )
+        if negative_prompt:
+            kwargs["negative_prompt"] = negative_prompt
+
+        result = pipe(**kwargs)
         return result.images[0]
 
     def generate_next_panel(
@@ -145,25 +179,27 @@ class ImageGenerator:
     ) -> Image.Image:
         """Generate the next panel using img2img from the previous panel."""
         pipe = get_pipeline(self.model_id, "i2i")
-        pipe = self._move_to_device(pipe)
 
         n_steps = steps or self.config["num_inference_steps"]
         cfg = guidance if guidance is not None else self.config["guidance_scale"]
         generator = torch.Generator(device=self.device).manual_seed(seed)
 
-        # Resize previous image to match model's expected resolution
+        # Resize previous image to match model resolution
         target_w, target_h = self.config["default_resolution"]
         prev_resized = previous_image.resize((target_w, target_h), Image.LANCZOS)
 
-        result = pipe(
+        kwargs = dict(
             prompt=prompt,
             image=prev_resized,
-            negative_prompt=negative_prompt if negative_prompt else None,
             num_inference_steps=n_steps,
             guidance_scale=cfg,
             strength=strength,
             generator=generator,
         )
+        if negative_prompt:
+            kwargs["negative_prompt"] = negative_prompt
+
+        result = pipe(**kwargs)
         return result.images[0]
 
     def generate_all_panels(
@@ -180,8 +216,8 @@ class ImageGenerator:
     ) -> List[Image.Image]:
         """
         Generate all panels in sequence.
-        Panel 1: text-to-image (or img2img from reference if provided)
-        Panel 2+: img2img from previous panel with controlled strength
+        Panel 1: text-to-image (or img2img from reference)
+        Panel 2+: img2img from previous panel
         """
         images = []
 
@@ -195,12 +231,11 @@ class ImageGenerator:
             try:
                 if i == 0:
                     if reference_image is not None:
-                        # Use reference image as base for first panel
                         img = self.generate_next_panel(
                             prompt=full_prompt,
                             previous_image=reference_image,
                             negative_prompt=neg_prompt,
-                            strength=max(strength, 0.5),  # Slightly higher for first panel from ref
+                            strength=max(strength, 0.5),
                             seed=seed,
                             steps=steps,
                             guidance=guidance,
@@ -230,7 +265,6 @@ class ImageGenerator:
 
             except Exception as e:
                 print(f"Error generating panel {i + 1}: {e}")
-                # Use previous image as fallback (or a blank image for panel 1)
                 if images:
                     images.append(images[-1].copy())
                 else:
@@ -253,10 +287,7 @@ class ImageGenerator:
         steps: Optional[int] = None,
         guidance: Optional[float] = None,
     ) -> Tuple[Image.Image, List[Image.Image]]:
-        """
-        Regenerate a single panel and optionally cascade changes forward.
-        Returns the new image and the updated image list.
-        """
+        """Regenerate a single panel."""
         prompt_data = prompts[panel_index]
         main_prompt = prompt_data.get("main_prompt", "")
         style = prompt_data.get("style_suffix", "")
@@ -266,17 +297,16 @@ class ImageGenerator:
 
         if panel_index == 0:
             new_img = self.generate_first_panel(
-                prompt=full_prompt, negative_prompt=neg_prompt, seed=seed, steps=steps, guidance=guidance
+                prompt=full_prompt, negative_prompt=neg_prompt,
+                seed=seed, steps=steps, guidance=guidance,
             )
         else:
             new_img = self.generate_next_panel(
                 prompt=full_prompt,
                 previous_image=existing_images[panel_index - 1],
                 negative_prompt=neg_prompt,
-                strength=strength,
-                seed=seed,
-                steps=steps,
-                guidance=guidance,
+                strength=strength, seed=seed,
+                steps=steps, guidance=guidance,
             )
 
         updated_images = existing_images.copy()